﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using NUnit.Framework;
using System.IO;
using SharpRBM.Core.Enumerations;

namespace OpenCL.Net.Tests
{
    [TestFixture]
    public class SharpRBMTests
    {
        private Cl.Context _context;
        private Cl.Device _device;

        [TestFixtureSetUp]
        public void Setup()
        {
            Cl.ErrorCode error;

            _device = (from device in
                           Cl.GetDeviceIDs(
                               (from platform in Cl.GetPlatformIDs(out error)
                                where Cl.GetPlatformInfo(platform, Cl.PlatformInfo.Name, out error).ToString() == "NVIDIA CUDA"
                                select platform).First(), Cl.DeviceType.Gpu, out error)
                       select device).First();

            _context = Cl.CreateContext(null, 1, new[] { _device }, null, IntPtr.Zero, out error);
        }

        [TestFixtureTearDown]
        public void Teardown()
        {
            _context.Dispose();
        }

        [Test]
        // Partially from OpenTK demo
        public void AddArrayAddsCorrectly()
        {
            const string correctSource = @"
                // Simple test; c[i] = a[i] + b[i]

                __kernel void add_array(__global float *a, __global float *b, __global float *c)
                {                   
                    
                    float x[100000];
                    unsigned int count= 100000;
                    for(int i=0;i<count;i++)
                    {
                        x[i]=1;
                    }
    
                    int xid = get_global_id(0);
                    c[xid] = a[xid] + b[xid];
                }
                
                __kernel void sub_array(__global float *a, __global float *b, __global float *c)
                {
                    int xid = get_global_id(0);
                    c[xid] = a[xid] - b[xid];
                }

                ";

            Cl.ErrorCode error;

            using (Cl.Program program = Cl.CreateProgramWithSource(_context, 1, new[] { correctSource }, null, out error))
            {
                error = Cl.BuildProgram(program, 1, new[] { _device }, string.Empty, null, IntPtr.Zero);
                if (error != Cl.ErrorCode.Success)
                {
                    char[] buffer = new char[9000];
                    Cl.InfoBuffer infoBuffer = new Cl.InfoBuffer(new IntPtr(9000));
                    IntPtr retSize;
                    Cl.GetProgramBuildInfo(program, _device, Cl.ProgramBuildInfo.Log, new IntPtr(9000), infoBuffer, out retSize);
                    Assert.Fail(infoBuffer.ToString());
                }
                Assert.AreEqual(Cl.ErrorCode.Success, error);
                Assert.AreEqual(Cl.BuildStatus.Success, Cl.GetProgramBuildInfo(program, _device, Cl.ProgramBuildInfo.Status, out error).CastTo<Cl.BuildStatus>());

                Cl.Kernel[] kernels = Cl.CreateKernelsInProgram(program, out error);
                Cl.Kernel kernel = kernels[0];

                const int cnBlockSize = 4;
                const int cnBlocks = 3;
                IntPtr cnDimension = new IntPtr(cnBlocks * cnBlockSize);

                // allocate host  vectors
                float[] A = new float[cnDimension.ToInt32()];
                float[] B = new float[cnDimension.ToInt32()];
                float[] C = new float[cnDimension.ToInt32()];

                // initialize host memory
                Random rand = new Random();
                for (int i = 0; i < A.Length; i++)
                {
                    A[i] = rand.Next() % 256;
                    B[i] = rand.Next() % 256;
                }

                Cl.Mem hDeviceMemA = Cl.CreateBuffer(_context, Cl.MemFlags.CopyHostPtr | Cl.MemFlags.ReadOnly, (IntPtr)(sizeof(float) * cnDimension.ToInt32()), A, out error);
                Assert.AreEqual(Cl.ErrorCode.Success, error);
                Cl.Mem hDeviceMemB = Cl.CreateBuffer(_context, Cl.MemFlags.CopyHostPtr | Cl.MemFlags.ReadOnly, (IntPtr)(sizeof(float) * cnDimension.ToInt32()), B, out error);
                Assert.AreEqual(Cl.ErrorCode.Success, error);
                Cl.Mem hDeviceMemC = Cl.CreateBuffer(_context, Cl.MemFlags.CopyHostPtr | Cl.MemFlags.WriteOnly, (IntPtr)(sizeof(float) * cnDimension.ToInt32()), IntPtr.Zero, out error);
                Assert.AreEqual(Cl.ErrorCode.Success, error);

                Cl.CommandQueue cmdQueue = Cl.CreateCommandQueue(_context, _device, (Cl.CommandQueueProperties)0, out error);

                Cl.Event clevent;

                int inPtrSize = 0;
                unsafe
                {
                    inPtrSize = sizeof(IntPtr);
                }

                // setup parameter values
                error = Cl.SetKernelArg(kernel, 0, new IntPtr(inPtrSize), hDeviceMemA);
                Assert.AreEqual(Cl.ErrorCode.Success, error);
                error = Cl.SetKernelArg(kernel, 1, new IntPtr(inPtrSize), hDeviceMemB);
                Assert.AreEqual(Cl.ErrorCode.Success, error);
                error = Cl.SetKernelArg(kernel, 2, new IntPtr(inPtrSize), hDeviceMemC);
                Assert.AreEqual(Cl.ErrorCode.Success, error);

                // write data from host to device
                error = Cl.EnqueueWriteBuffer(cmdQueue, hDeviceMemA, Cl.Bool.True, IntPtr.Zero,
                    new IntPtr(cnDimension.ToInt32() * sizeof(float)),
                    A, 0, null, out clevent);
                Assert.AreEqual(Cl.ErrorCode.Success, error);
                error = Cl.EnqueueWriteBuffer(cmdQueue, hDeviceMemB, Cl.Bool.True, IntPtr.Zero,
                    new IntPtr(cnDimension.ToInt32() * sizeof(float)),
                    B, 0, null, out clevent);
                Assert.AreEqual(Cl.ErrorCode.Success, error);

                // execute kernel
                error = Cl.EnqueueNDRangeKernel(cmdQueue, kernel, 1, null, new IntPtr[] { cnDimension }, null, 0, null, out clevent);
                Assert.AreEqual(Cl.ErrorCode.Success, error, error.ToString());

                // copy results from device back to host
                IntPtr event_handle = IntPtr.Zero;

                error = Cl.EnqueueReadBuffer(
                    cmdQueue, hDeviceMemC, Cl.Bool.True, IntPtr.Zero,
                    new IntPtr(cnDimension.ToInt32() * sizeof(float)),
                    C, 0, null, out clevent);
                Assert.AreEqual(Cl.ErrorCode.Success, error, error.ToString());


                for (int i = 0; i < A.Length; i++)
                {
                    Assert.That(A[i] + B[i], Is.EqualTo(C[i]));
                }

                Cl.Finish(cmdQueue);
                Cl.ReleaseMemObject(hDeviceMemA);
                Cl.ReleaseMemObject(hDeviceMemB);
                Cl.ReleaseMemObject(hDeviceMemC);
            }
        }

        [Test]
        // Partially from OpenTK demo
        public void MultiMethodTestWorksCorrectly()
        {
            const string correctSource = @"
                // Two step test
               
                __kernel void sub_array(__global float *a, __global float *b, __global float *c)
                {
                    int xid = get_global_id(0);
                    c[xid] = a[xid] - b[xid];
                }
                
                __kernel void add_array(__global float *a, __global float *b, __global float *c)
                {
                    int xid = get_global_id(0);
                    c[xid] = a[xid] + b[xid];
                }

                __kernel void call_sub_array(__global float *a, __global float *b, __global float *c)
                {    
                    sub_array(a,b,c);
                }
                ";

            Cl.ErrorCode error;

            using (Cl.Program program = Cl.CreateProgramWithSource(_context, 1, new[] { correctSource }, null, out error))
            {
                Assert.AreEqual(error, Cl.ErrorCode.Success);
                error = Cl.BuildProgram(program, 1, new[] { _device }, string.Empty, null, IntPtr.Zero);
                Assert.AreEqual(Cl.ErrorCode.Success, error);
                Assert.AreEqual(Cl.GetProgramBuildInfo(program, _device, Cl.ProgramBuildInfo.Status, out error).CastTo<Cl.BuildStatus>(), Cl.BuildStatus.Success);
                Cl.Kernel kernel = Cl.CreateKernel(program, "call_sub_array", out error);

                const int cnBlockSize = 4;
                const int cnBlocks = 3;
                IntPtr cnDimension = new IntPtr(cnBlocks * cnBlockSize);

                // allocate host  vectors
                float[] A = new float[cnDimension.ToInt32()];
                float[] B = new float[cnDimension.ToInt32()];
                float[] C = new float[cnDimension.ToInt32()];

                // initialize host memory
                Random rand = new Random();
                for (int i = 0; i < A.Length; i++)
                {
                    A[i] = rand.Next() % 256;
                    B[i] = rand.Next() % 256;
                }

                Cl.Mem hDeviceMemA = Cl.CreateBuffer(_context, Cl.MemFlags.CopyHostPtr | Cl.MemFlags.ReadOnly, (IntPtr)(sizeof(float) * cnDimension.ToInt32()), A, out error);
                Assert.AreEqual(Cl.ErrorCode.Success, error);
                Cl.Mem hDeviceMemB = Cl.CreateBuffer(_context, Cl.MemFlags.CopyHostPtr | Cl.MemFlags.ReadOnly, (IntPtr)(sizeof(float) * cnDimension.ToInt32()), B, out error);
                Assert.AreEqual(Cl.ErrorCode.Success, error);
                Cl.Mem hDeviceMemC = Cl.CreateBuffer(_context, Cl.MemFlags.CopyHostPtr | Cl.MemFlags.WriteOnly, (IntPtr)(sizeof(float) * cnDimension.ToInt32()), IntPtr.Zero, out error);
                Assert.AreEqual(Cl.ErrorCode.Success, error);

                Cl.CommandQueue cmdQueue = Cl.CreateCommandQueue(_context, _device, (Cl.CommandQueueProperties)0, out error);

                Cl.Event clevent;

                int inPtrSize = 0;
                unsafe
                {
                    inPtrSize = sizeof(IntPtr);
                }

                // setup parameter values
                error = Cl.SetKernelArg(kernel, 0, new IntPtr(inPtrSize), hDeviceMemA);
                Assert.AreEqual(Cl.ErrorCode.Success, error);
                error = Cl.SetKernelArg(kernel, 1, new IntPtr(inPtrSize), hDeviceMemB);
                Assert.AreEqual(Cl.ErrorCode.Success, error);
                error = Cl.SetKernelArg(kernel, 2, new IntPtr(inPtrSize), hDeviceMemC);
                Assert.AreEqual(Cl.ErrorCode.Success, error);

                // write data from host to device
                error = Cl.EnqueueWriteBuffer(cmdQueue, hDeviceMemA, Cl.Bool.True, IntPtr.Zero,
                    new IntPtr(cnDimension.ToInt32() * sizeof(float)),
                    A, 0, null, out clevent);
                Assert.AreEqual(Cl.ErrorCode.Success, error);
                error = Cl.EnqueueWriteBuffer(cmdQueue, hDeviceMemB, Cl.Bool.True, IntPtr.Zero,
                    new IntPtr(cnDimension.ToInt32() * sizeof(float)),
                    B, 0, null, out clevent);
                Assert.AreEqual(Cl.ErrorCode.Success, error);

                // execute kernel
                error = Cl.EnqueueNDRangeKernel(cmdQueue, kernel, 1, null, new IntPtr[] { cnDimension }, null, 0, null, out clevent);
                Assert.AreEqual(Cl.ErrorCode.Success, error, error.ToString());

                // copy results from device back to host
                IntPtr event_handle = IntPtr.Zero;

                error = Cl.EnqueueReadBuffer(
                    cmdQueue, hDeviceMemC, Cl.Bool.True, IntPtr.Zero,
                    new IntPtr(cnDimension.ToInt32() * sizeof(float)),
                    C, 0, null, out clevent);
                Assert.AreEqual(Cl.ErrorCode.Success, error, error.ToString());

                for (int i = 0; i < A.Length; i++)
                {
                    Assert.That(C[i], Is.EqualTo(A[i] - B[i]));
                }

                Cl.Finish(cmdQueue);
                Cl.ReleaseMemObject(hDeviceMemA);
                Cl.ReleaseMemObject(hDeviceMemB);
                Cl.ReleaseMemObject(hDeviceMemC);
            }
        }

        [Test]
        public void TestTrainRBM()
        {
            string correctSource = File.ReadAllText(@"..\..\TrainRBM.c");
            correctSource =
                correctSource
                    .Replace("<UpperCount>", "1001")
                    .Replace("<LowerCount>", "1002");

            Cl.ErrorCode error;

            using (Cl.Program program = Cl.CreateProgramWithSource(_context, 1, new[] { correctSource }, null, out error))
            {
                error = Cl.BuildProgram(program, 1, new[] { _device }, string.Empty, null, IntPtr.Zero);
                if (error != Cl.ErrorCode.Success)
                {
                    char[] buffer = new char[9000];
                    Cl.InfoBuffer infoBuffer = new Cl.InfoBuffer(new IntPtr(9000));
                    IntPtr retSize;
                    Cl.GetProgramBuildInfo(program, _device, Cl.ProgramBuildInfo.Log, new IntPtr(9000), infoBuffer, out retSize);
                    Assert.Fail(infoBuffer.ToString());
                }

                Assert.AreEqual(error, Cl.ErrorCode.Success);
                error = Cl.BuildProgram(program, 1, new[] { _device }, string.Empty, null, IntPtr.Zero);
                Assert.AreEqual(Cl.ErrorCode.Success, error);
                Assert.AreEqual(Cl.GetProgramBuildInfo(program, _device, Cl.ProgramBuildInfo.Status, out error).CastTo<Cl.BuildStatus>(), Cl.BuildStatus.Success);
                Cl.Kernel kernel = Cl.CreateKernel(program, "TrainOnSingleCase", out error);
                Assert.AreEqual(Cl.ErrorCode.Success, error);

            }
        }

        [Test]
        public void HasMassiveMemoryLeak()
        {
            string correctSource =
@"__kernel void ActivateUpperToLower(
	__global float *lowerValues, 
	__global float *upperValues, 
	__global float *weights)
{
	int lower = get_global_id(0);
}";

            Cl.ErrorCode error;

            using (Cl.Program program = Cl.CreateProgramWithSource(_context, 1, new[] { correctSource }, null, out error))
            {
                error = Cl.BuildProgram(program, 1, new[] { _device }, string.Empty, null, IntPtr.Zero);
                if (error != Cl.ErrorCode.Success)
                {
                    char[] buffer = new char[9000];
                    Cl.InfoBuffer infoBuffer = new Cl.InfoBuffer(new IntPtr(9000));
                    IntPtr retSize;
                    Cl.GetProgramBuildInfo(program, _device, Cl.ProgramBuildInfo.Log, new IntPtr(9000), infoBuffer, out retSize);
                    Assert.Fail(infoBuffer.ToString());
                }

                Assert.AreEqual(error, Cl.ErrorCode.Success);
                error = Cl.BuildProgram(program, 1, new[] { _device }, string.Empty, null, IntPtr.Zero);
                Assert.AreEqual(Cl.ErrorCode.Success, error);
                Assert.AreEqual(Cl.GetProgramBuildInfo(program, _device, Cl.ProgramBuildInfo.Status, out error).CastTo<Cl.BuildStatus>(), Cl.BuildStatus.Success);
                Cl.Kernel kernel = Cl.CreateKernel(program, "ActivateUpperToLower", out error);
                Assert.AreEqual(Cl.ErrorCode.Success, error);

                Cl.CommandQueue commandQueue = Cl.CreateCommandQueue(_context, _device, (Cl.CommandQueueProperties)0, out error);
                float[] array = new float[10];
                Cl.Mem mem = Cl.CreateBuffer(_context, Cl.MemFlags.CopyHostPtr | Cl.MemFlags.ReadWrite, (IntPtr)(sizeof(float) * array.Length), array, out error);

                for (int i = 0; i < 1000000; i++)
                {
                    Cl.Event clevent;
                    IntPtr event_handle = IntPtr.Zero;
                    error = Cl.EnqueueReadBuffer(commandQueue, mem, Cl.Bool.True, IntPtr.Zero, new IntPtr(array.Length * sizeof(float)), array, 0, null, out clevent);
                    // **************************************
                    // THIS IS WHAT CAUSES THE MEMORY LEAK!!!
                    // **************************************
                    // clevent.Dispose();
                }
            }
        }

        [Test]
        public void NoMemoryLeak()
        {
            string correctSource =
@"__kernel void ActivateUpperToLower(
	__global float *lowerValues, 
	__global float *upperValues, 
	__global float *weights)
{
	int lower = get_global_id(0);
}";

            Cl.ErrorCode error;

            using (Cl.Program program = Cl.CreateProgramWithSource(_context, 1, new[] { correctSource }, null, out error))
            {
                error = Cl.BuildProgram(program, 1, new[] { _device }, string.Empty, null, IntPtr.Zero);
                if (error != Cl.ErrorCode.Success)
                {
                    char[] buffer = new char[9000];
                    Cl.InfoBuffer infoBuffer = new Cl.InfoBuffer(new IntPtr(9000));
                    IntPtr retSize;
                    Cl.GetProgramBuildInfo(program, _device, Cl.ProgramBuildInfo.Log, new IntPtr(9000), infoBuffer, out retSize);
                    Assert.Fail(infoBuffer.ToString());
                }

                Assert.AreEqual(error, Cl.ErrorCode.Success);
                error = Cl.BuildProgram(program, 1, new[] { _device }, string.Empty, null, IntPtr.Zero);
                Assert.AreEqual(Cl.ErrorCode.Success, error);
                Assert.AreEqual(Cl.GetProgramBuildInfo(program, _device, Cl.ProgramBuildInfo.Status, out error).CastTo<Cl.BuildStatus>(), Cl.BuildStatus.Success);
                Cl.Kernel kernel = Cl.CreateKernel(program, "ActivateUpperToLower", out error);
                Assert.AreEqual(Cl.ErrorCode.Success, error);

                Cl.CommandQueue commandQueue = Cl.CreateCommandQueue(_context, _device, (Cl.CommandQueueProperties)0, out error);
                if (error != Cl.ErrorCode.Success)
                {
                    throw new InvalidProgramException(string.Format("Unable to CreateCommandQueue: {0}!", error));
                }

                float[] array = new float[10];
                Cl.Mem mem = Cl.CreateBuffer(_context, Cl.MemFlags.CopyHostPtr | Cl.MemFlags.ReadWrite, (IntPtr)(sizeof(float) * array.Length), array, out error);

                for (int i = 0; i < 1000000; i++)
                {
                    Cl.Event clevent;
                    IntPtr event_handle = IntPtr.Zero;
                    error = Cl.EnqueueReadBuffer(commandQueue, mem, Cl.Bool.True, IntPtr.Zero, new IntPtr(array.Length * sizeof(float)), array, 0, null, out clevent);
                    // *****************************************
                    // THIS IS WHAT PREVENTS THE MEMORY LEAK!!!
                    // *****************************************
                    clevent.Dispose();
                }
                if (error != Cl.ErrorCode.Success)
                {
                    throw new InvalidOperationException("EnqueueReadBuffer failed for _weightsMem:" + error);
                }
                kernel.SetKernelArgFloat(0, mem);
            }
        }
    }
}
